from flask import Flask, request, jsonify

from news_classification.glda import get_glda_infer_engine, infer_doc_topics_by_glda

app = Flask(__name__)

model_fp = 'model/1M/glda-128-topics/1M-news-128-topics-lda.model'
topic_num = 128
topics_assign = {
    57: '业务增减', 99: '业绩增长', 81: '行业市场',

    71: '会议公告', 114: '上市发行', 29: '债务相关', 100: '股权质押',

    111: '新冠药物', 3: '疫情影响', 11: '新冠肺炎', 23: '救援武汉', 105: '疫情捐赠',
    26: '疫情防控', 31: '复工复产', 48: '疫情防控', 85: '疫情防控', 123: '疫情防控',

    116: '税务',
    112: '进出口',
    8: '银行金融', 77: '融资贷款', 98: '融资募集', 82: '基金证券', 49: '保险',

    60: '智能手机', 94: '智能手机', 42: '5G通讯',
    118: '影视娱乐',
    74: '电商零售',
    5: '医药', 86: '医疗', 66: '医疗器械',
    0: '航空、无人机', 91: '航空动力',
    108: '芯片半导体',
    6: '智能技术', 117: '区块链',
    16: '农业农产', 38: '养殖业',
    102: '物流', 20: '交通运输',
    1: '新能源汽车', 88: '自动驾驶',
    24: '电力',
    32: '在线教育',
    64: '白酒',
    41: '文旅休闲',
    47: '能源', 63: '清洁能源',
    106: '家电',
    109: '共享出行',
    68: '学校教育',
    75: '体育竞技',
    84: '地产',
    7: '生活服务',
    18: '商场餐饮',
    27: '租赁',
    28: '住房',
    37: '品牌时尚',

    # 39: '产业投资', 52: '科技投资',
    # 69: '产业发展', 55: '经济发展',

    125: '资产变动',
    78: '企业高管',
    4: '招标采购', 45: '招投标',
    113: '生产供应',
    97: '人才就业',
    103: '知识产权',

    2: '扶贫攻坚',
    9: '员工工会',
    12: '政府补助',
    # 13: '工程建设', 25: '工程建设',

    14: '国际海外',
    17: '山东', 30: '浙江', 22: '长三角', 54: '广东', 67: '江苏', 89: '四川', 107: '东北', 96: '安徽',
    # 19: '社会组织',

    50: '市场监管', 80: '调查监督',
    10: '环保污染',
    90: '食品安全', 95: '食品安全',
    93: '消费投诉', 119: '违法处罚',
    70: '公安执法',
    58: '质量管控',
}
inf_engine = get_glda_infer_engine(model_fp)


def get_assign_topic(topic_id):
    try:
        return topics_assign[topic_id]
    except KeyError:
        return 'Topic %d' % topic_id


@app.route('/infer_doc_topics', methods=['POST'])
def infer_doc_topics():
    doc = request.form['doc']
    topic_dis = infer_doc_topics_by_glda(doc, inf_engine)
    result = []
    for topic_id, prob in topic_dis:
        if prob < 0.1:
            continue
        result.append(
            {
                'topic_id': topic_id,
                'topic_name': get_assign_topic(topic_id),
                'topic_weight': prob
            }
        )
    return jsonify(result)


@app.route('/fetch_doc_topic_vector', methods=['POST'])
def fetch_doc_topic_vector():
    doc = request.form['doc']
    topic_dis = infer_doc_topics_by_glda(doc, inf_engine)
    vector = [0] * topic_num
    for topic_id, prob in topic_dis:
        vector[topic_id] = prob
    return jsonify(vector)


model_fp2 = 'model/1M/glda-512-topics/1M-news-512-topics-lda.model'
topic_num2 = 512
inf_engine2 = get_glda_infer_engine(model_fp2)


@app.route('/fetch_doc_512topic_vector', methods=['POST'])
def fetch_doc_512topic_vector():
    doc = request.form['doc']
    topic_dis = infer_doc_topics_by_glda(doc, inf_engine2)
    vector = [0] * topic_num2
    for topic_id, prob in topic_dis:
        vector[topic_id] = prob
    return jsonify(vector)


model_fp3 = 'model/3M/1000-topics/3M-news-1000-topics-gensim-lda.model'
topic_num3 = 1000
inf_engine3 = get_glda_infer_engine(model_fp3)


@app.route('/fetch_doc_1000topic_vector', methods=['POST'])
def fetch_doc_1000topic_vector():
    doc = request.form['doc']
    topic_dis = infer_doc_topics_by_glda(doc, inf_engine3)
    vector = [0] * topic_num3
    for topic_id, prob in topic_dis:
        vector[topic_id] = prob
    return jsonify(vector)
